/*
 * Alloy Analyzer
 * Copyright (c) 2005 Massachusetts Institute of Technology
 *
 * The Alloy Analyzer is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * The Alloy Analyzer is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with the Alloy Analyzer; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */
package alloy.gui;

import java.awt.Color;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;

import javax.swing.text.DefaultHighlighter;
import javax.swing.text.Highlighter;

import alloy.api.LiftedUnsatCore;
import alloy.ast.ASTDepthFirstVisitor;
import alloy.ast.Formula;
import alloy.ast.FormulaSeq;
import alloy.ast.Formulas;
import alloy.ast.Location;
import alloy.ast.Node;
import alloy.ast.TreeNode;
import alloy.bool.UcoreProofProcessor;
import alloy.util.Dbg;
import alloy.util.Params;
import alloy.util.SetOfLong;

/**
 * @author edmond
 * @author rseater
 * 
 * The LiftedUnsatCoreHighlightVisitor is constructed with a
 * {@link alloy.api.LiftedUnsatCore}. It traverses an AST tree to determine
 * which nodes are part of the unsat core and then uses three different shades
 * of highlighting to visualize the unsat core in the text area.
 * 
 * The {@link alloy.api.LiftedUnsatCore}provides an abstraction to query AST
 * nodes and see whether they are in the unsat core. The boolean translation of
 * a desugared AST is a tree of boolean formulas. The unsat core shares the root
 * and forms a subtree within this larger tree. Due to sharing, nodes deeper in
 * the tree may appear to be floating and detached from the unsat core tree. To
 * mitigate that case, this visitor traverses the tree depth-first but stops
 * traversing a particular branch when it finds a node not in the unsat core.
 * 
 * The LiftedUnsatCoreHighlightVisitor works in three phases. In the first
 * phase, it walks the AST, classifying nodes that are in the unsat core and
 * gathering the names of the different files used. Then the client (i.e.
 * {@link alloy.gui.NewAlloyGUI}) maps each used file to an associated text
 * area. This step is to get around the fact that attempting to highlight a node
 * in a file other than the one open in the text area causes that new file to be
 * opened in the same text area, clobbering the previously open file. Thus, each
 * file has its own associated text area. Finally, the
 * LiftedUnsatCoreHighlightVisitor highlights each node in the text area
 * associated with its location's file.
 * 
 * To see how to check which quantified variables are in the unsat core, look at
 * alloy.gui.SimpleASTSwingVisitor in alloy2-1.
 * 
 * Highlighting seems to only work if more specific locations are highlighted
 * before more general locations. If location A is contained in location B,
 * highlighting B in one color then highlighting A in a different color seems to
 * only do the first hightlight. Highlighting location A and then B works fine.
 */
public class LiftedUnsatCoreHighlightVisitor extends ASTDepthFirstVisitor {

	/** the {@link LiftedUnsatCore}computed by {@link AlloyRunner} */
	private LiftedUnsatCore _core;

	/** {@link Map}from files to text areas for determining where to highlight */
	private Map /* <File->LocatableTextArea> */_filesToTextAreas;

	/*
	 * The LiftedUnsatCoreHighlightVisitor visits AST nodes and assigns those in
	 * the unsat core into three different categories. Undert the syntatic
	 * method of highlighting, these categories are highlighted in three
	 * different shades.
	 * 
	 * Alloy facts, predicates, and functions are represented in the AST tree as
	 * Formulas and are given the lightest shade of highlighting. Individual
	 * formulas in these paragraphs, represented as Formula nodes, are given a
	 * medium shade of highlighting.
	 * 
	 * Separate tokens in the individual formulas are given the darkest shade of
	 * highlighting
	 */
	/**
	 * {@link Set}of AST Formula nodes collected by visitor. Under syntactic
	 * highlighting, their locations to be highlighted in lightest shade.
	 */
	private Set /* <Formulas> */_paragraphsInCore;

	/**
	 * {@link Set}of AST Formula nodes collected by visitor. Under syntactic
	 * highlighting, their locations to be highlighted in medium shade.
	 */
	private Set /* <Formula> */_formulasInCore;

	/**
	 * {@link Set}of other AST tokens collected by visitor. Under syntactic
	 * highlighting, their locations to be highlighted in darkest shade.
	 */
	private Set /* <Node> */_tokensInCore;

	/**
	 * {@link Map}from nodes to the frequency of their corresponding literals
	 * in the core. It is filled in by examining rawLiteralCounts during the
	 * visit() traversal of the AST
	 */
	private Map /* Node to Integer */_templateCountsForNodes;

	/**
	 * {@link Set}of all invocation locations of functions.
	 */
	private Set /* <Location> */_invocationLocations;

	/**
	 * Constuctor
	 * 
	 * @param core_
	 */
	public LiftedUnsatCoreHighlightVisitor(final LiftedUnsatCore core_) {
		Dbg.chk(core_);
		_core = core_;

		_paragraphsInCore = new HashSet();
		_formulasInCore = new HashSet();
		_tokensInCore = new HashSet();
		_templateCountsForNodes = new HashMap();

		_invocationLocations = new HashSet();
		_filesToTextAreas = new HashMap();
	}

	/**
	 * Public factory method to create LiftedUnsatCoreHighlightVisitor and to
	 * highlight the unsat core in the text area
	 * 
	 * @param core_
	 *            LiftedUnsatCore
	 * @param formula_
	 *            the AST root of the formula to analyze
	 * @param translator_
	 *            record of translations from AST nodes to boolean formulas
	 * @param textArea_
	 *            text area with the GUI
	 */
	public static LiftedUnsatCoreHighlightVisitor constructVisitor(final LiftedUnsatCore core_, final Formula formula_) {
		final LiftedUnsatCoreHighlightVisitor visitor = new LiftedUnsatCoreHighlightVisitor(core_);
		formula_.applyVisitor(visitor);
		return visitor;
	}

	/**
	 * Adds a mapping from the filename_ to textArea_
	 * 
	 * @param filename_
	 * @param textArea_
	 * @throws NullPointerException
	 *             if filename_ or textArea_ is null
	 * @throws IllegalArgumentException
	 *             if filename_ is already mapped
	 */
	public void mapFileToTextArea(final String filename_, final LocatableTextArea textArea_) {
		if (filename_ == null || textArea_ == null)
			throw new NullPointerException();
		if (_filesToTextAreas.containsKey(filename_) && _filesToTextAreas.get(filename_) != null)
			throw new IllegalArgumentException("File " + filename_ + " already mapped.");
		_filesToTextAreas.put(filename_, textArea_);
	}

	/**
	 * Returns the Set[String] of filenames for the AST.
	 */
	public Set getFileNames() {
		return Collections.unmodifiableSet(_filesToTextAreas.keySet());
	}

	/**
	 * Highlights the node's location with the given painter in the text area
	 * that is mapped to by the node's filename. If the node has no associated
	 * file, does nothing.
	 * 
	 * @throws IllegalStateException
	 *             if node's associated file is not mapped to a text area
	 * @param node_
	 *            unsat core node to be highlighted
	 * @param painter_
	 *            the painter to highlight with
	 */
	private void highlightNode(final Node node_, final Highlighter.HighlightPainter painter_) {
		final String filename = getLocationFileName(node_);
		if (filename == null)
			return;
		final LocatableTextArea textArea = (LocatableTextArea) _filesToTextAreas.get(filename);
		if (textArea == null) {
			throw new IllegalStateException("Filename " + filename
					+ " for node " + node_
					+ " not mapped to a text area");
		}

		textArea.addLoc(node_.getLocation(), painter_); //this adds the highlighting
	}

	/**
	 * highlightNode(n) is like the highlightNode(n,p) method, except that
	 * it selects its own color painter. The color is selected based on
	 * annotations on the node about its frequency in the core.
	 * 
	 * Highlights the node's location with the given painter in the text area
	 * that is mapped to by the node's filename. If the node has no associated
	 * file, does nothing.
	 * 
	 * @throws IllegalStateException
	 *             if node's associated file is not mapped to a text area
	 * @param node_
	 *            unsat core node to be highlighted
	 */
	private void highlightNode(final Node node_) {
		String filename = getLocationFileName(node_);
		if (filename == null)
			return;
		LocatableTextArea textArea = (LocatableTextArea) _filesToTextAreas.get(filename);
		if (textArea == null) {
			throw new IllegalStateException("Filename " + filename
					+ " for node " + node_ + " not mapped to a text area");
		}

		//assert !(Params.glob.getParam("DEVEL", "unsatCoreStyle").equals("Syntactic"));
		final Highlighter.HighlightPainter painter;
		if (Params.glob.getParam("DEVEL", "unsatCoreStyle").equals("Uniform")) {
			painter = new DefaultHighlighter.DefaultHighlightPainter(new Color(150, 150, 175)); //bluish gray
		} else {
			//assert Params.glob.getParam("DEVEL", "unsatCoreStyle").equals("HitCount");
			if (Params.glob.getParam("DEVEL", "hitCountStyle").equals("literalCountNoClauseDuplication")) {
				final int literalCount = ((Integer) _templateCountsForNodes.get(node_)).intValue();
				final int emphasis = Math.min(205, Math.max(50, 255 - (25 + 50 * literalCount))); //50 to 205
				painter = new DefaultHighlighter.DefaultHighlightPainter(new Color(emphasis + 15, emphasis - 20, emphasis + 45));
				//purplish-gray: darker is more emphatic
			} else {
				//assert !(Params.glob.getParam("DEVEL", "hitCountStyle").equals("literalCountNoClauseDuplication"))
				painter = new DefaultHighlighter.DefaultHighlightPainter(new Color(100, 100, 100)); //light gray
			}
		}
		textArea.addLoc(node_.getLocation(), painter); //this adds the highlighting
	}

	/**
	 * @param node_
	 *            what to unhighlight
	 * 
	 * @throws IllegalStateException
	 *             if node's associated file is not mapped to a text area
	 * @param node_
	 *            unsat core node to be highlighted
	 * @param painter_
	 *            the painter to highlight with
	 */
	private void unhighlightNode(final Node node_) {
		final String filename = getLocationFileName(node_);
		if (filename == null)
			return;
		final LocatableTextArea textArea = (LocatableTextArea) _filesToTextAreas.get(filename);
		if (textArea == null) {
			throw new IllegalStateException("Filename " + filename + " for node " + node_ + " not mapped to a text area");
		}

		textArea.clearLocs(); //this eliminates the highlighting
	}

	/**
	 * @throws IllegalStateException
	 *             if node's associated file is not mapped to a text area
	 * @param location_
	 *            what to highlight
	 * @param painter_
	 *            how to highlight it
	 */
	private void highlightLocation(final Location location_, final Highlighter.HighlightPainter painter_) {
		if (location_ == null || location_.getFileName() == null || location_.getFileName() == Location.NOT_FROM_FILE)
			return;
		final String filename = location_.getFileName();

		final LocatableTextArea textArea = (LocatableTextArea) _filesToTextAreas.get(filename);
		if (textArea == null) {
			throw new IllegalStateException("Filename " + filename
					+ " for location " + location_
					+ " not mapped to a text area");
		}

		textArea.addLoc(location_, painter_); //this adds the highlighting
	}

	/**
	 * @throws IllegalStateException
	 *             if node's associated file is not mapped to a text area
	 * @param location_
	 *            what to highlight
	 * 
	 * highlightNode(n) is like the other highlightNode(n,p) method, except that
	 * it selects its own color painter. The color is selected based on
	 * annotations on the node about its frequency in the core.
	 */
	private void highlightLocation(final Location location_) {
		if (location_ == null || location_.getFileName() == null
				|| location_.getFileName() == Location.NOT_FROM_FILE)
			return;
		final String filename = location_.getFileName();

		final LocatableTextArea textArea = (LocatableTextArea) _filesToTextAreas
				.get(filename);
		if (textArea == null) {
			throw new IllegalStateException("Filename " + filename
					+ " for location " + location_
					+ " not mapped to a text area");
		}

		//does this actually get called? -rseater
		//assert false
		final Highlighter.HighlightPainter painter = new DefaultHighlighter.DefaultHighlightPainter(
				new Color(50, 50, 50)); //very dark gray
		textArea.addLoc(location_, painter); //this adds the highlighting
	}

	/**
	 * @param location_
	 *            what to unhighlight
	 */
	private void unhighlightLocation(final Location location_) {
		if (location_ == null || location_.getFileName() == null || location_.getFileName() == Location.NOT_FROM_FILE)
			return;
		final String filename = location_.getFileName();

		final LocatableTextArea textArea = (LocatableTextArea) _filesToTextAreas.get(filename);
		if (textArea == null) {
			throw new IllegalStateException("Filename " + filename
					+ " for location " + location_
					+ " not mapped to a text area");
		}

		textArea.clearLocs(); //this eliminates the highlighting
	}

	/**
	 * Returns true if loc_ is subsumed by any location occupied by a Node in
	 * nodes_. Otherwise, returns false.
	 * 
	 * @param loc_
	 * @param locations_
	 *            Set[Location]
	 */
	private boolean subsumedBy(final Location loc_, final Set nodes_) {
		for (Iterator i = nodes_.iterator(); i.hasNext();) {
			final Node node = (Node) i.next();
			if (node.getLocation().subsumes(loc_))
				return true;
		}
		return false;
	}

	/**
	 * Returns the set of invocation locations s.t. each invocation location is
	 * subsumed by some paragraph in the core.
	 */
	private Set filterInvocationLocations() {
		Set filteredInvocationLocations = new HashSet();
		for (Iterator i = _invocationLocations.iterator(); i.hasNext();) {
			final Location loc = (Location) i.next();
			if (subsumedBy(loc, _paragraphsInCore)) {
				filteredInvocationLocations.add(loc);
				continue;
			}
		}
		return filteredInvocationLocations;
	}

	/**
	 * Tells if a node is in the unsat core and thus should be highlighted.
	 * 
	 * @param node_
	 *            the AST node
	 * @return true if node should be highlighted.
	 */
	private boolean shouldHighlightNode(final Node node_) {
		return _core.isInCore(node_);
	}

	/**
	 * Tells how often a node's literal is in the unsat core, and thus how
	 * emphatically it should be highlighted.
	 * 
	 * @param node_
	 *            the AST node
	 * @return how many times that node's literal appears in the core
	 */
	private Integer templateCountForNode(final Node node_) {
		final SetOfLong templates = _core.nodeTemplates(node_);
		if (templates == null)
			return new Integer(0);
		int totalNumberOfOccurences = 0;
		//iterate over templates and count each by asking ucore proof processor
		for (Iterator i = templates.toSet().iterator(); i.hasNext();) {
			final long currentTemplate = ((Long) i.next()).longValue();
			//assert currentTemplate <= Integer.MAX_VALUE;
			totalNumberOfOccurences += UcoreProofProcessor.getTemplateCount((int) currentTemplate);
		}
		return new Integer(totalNumberOfOccurences);
		//final int count = templates.size();
		//return new Integer(count);
	}

	/**
	 * @effect highlight the model pane based on the unsat core data and the
	 *         preferences
	 */
	public void showLiftedUnsatCore() {
		// check that all filenames have been mapped to text areas
		Dbg.chk(!_filesToTextAreas.values().contains(null));
		//assert (Params.glob.defined("DEVEL", "unsatCoreStyle"));
		if (Params.glob.getParam("DEVEL", "unsatCoreStyle").equals("Syntactic")) {
			//structure based highlighting
			final Highlighter.HighlightPainter _darkPainter =
				new DefaultHighlighter.DefaultHighlightPainter(new Color(120, 60, 255)); //dark purple
			for (Iterator k = _tokensInCore.iterator(); k.hasNext();)
				highlightNode((Node) k.next(), _darkPainter);
			for (Iterator k = filterInvocationLocations().iterator(); k.hasNext();)
				highlightLocation((Location) k.next(), _darkPainter);

			final Highlighter.HighlightPainter _mediumPainter =
				new DefaultHighlighter.DefaultHighlightPainter(new Color(170, 100, 255)); //medium purple
			for (Iterator j = _formulasInCore.iterator(); j.hasNext();)
				highlightNode((Node) j.next(), _mediumPainter);

			final Highlighter.HighlightPainter _lightPainter =
				new DefaultHighlighter.DefaultHighlightPainter(new Color(200, 160, 255)); //light purple
			for (Iterator i = _paragraphsInCore.iterator(); i.hasNext();)
				highlightNode((Node) i.next(), _lightPainter);
		} else {
			//assert (Params.glob.getParam("DEVEL",
			// "unsatCoreStyle").equals("HitCount"))
			//    || (Params.glob.getParam("DEVEL",
			// "unsatCoreStyle").equals("Uniform"));
			//TODO how to highlight just lines?
			for (Iterator k = _tokensInCore.iterator(); k.hasNext();)
				highlightNode((Node) k.next());
			for (Iterator k = filterInvocationLocations().iterator(); k.hasNext();)
				highlightLocation((Location) k.next());//rseater thinks this call doesn't matter
			for (Iterator j = _formulasInCore.iterator(); j.hasNext();)
				highlightNode((Node) j.next());
			for (Iterator i = _paragraphsInCore.iterator(); i.hasNext();)
				highlightNode((Node) i.next());
		}
	}

	/**
	 * @effect removes highlighting
	 */
	public void hideLiftedUnsatCore() {
		Dbg.chk(!_filesToTextAreas.values().contains(null));
		Dbg.info("Erasing old unsat core.");

		for (Iterator k = _tokensInCore.iterator(); k.hasNext();)
			unhighlightNode((Node) k.next());
		for (Iterator k = filterInvocationLocations().iterator(); k.hasNext();)
			unhighlightLocation((Location) k.next());
		for (Iterator j = _formulasInCore.iterator(); j.hasNext();)
			unhighlightNode((Node) j.next());
		for (Iterator i = _paragraphsInCore.iterator(); i.hasNext();)
			unhighlightNode((Node) i.next());
	}

	/**
	 * Performs a depth-first traversal of the AST tree, short circuiting a
	 * branch if a node not in the unsat core is encountered. Adds the node to
	 * the appropriate category based on type.
	 * 
	 * @param node_
	 *            the AST node
	 */
	public void visit(final Node node_) {
		if (node_ == null)
			throw new NullPointerException();

		final String filename = getLocationFileName(node_);

		// gather up the location information
		if (filename != null && !_filesToTextAreas.containsKey(filename))
			_filesToTextAreas.put(filename, null);

		// only visit children if current node is in unsat core
		if (shouldHighlightNode(node_)) {
			//syntactic/structural highlighting: add to appropriate set based on type
			if (node_ instanceof Formulas) {
				_paragraphsInCore.add(node_);
				List invocationLocations = ((TreeNode) node_)
						.getInvocationLocations();
				if (invocationLocations.size() != 0)
					_invocationLocations.addAll(invocationLocations);
			} else if (node_ instanceof Formula) {
				_formulasInCore.add(node_);
			} else {//instance of a token
				_tokensInCore.add(node_);
			}
			//add node to hit count mapping for later use by hitcount highlighting
			if (_templateCountsForNodes == null)
				_templateCountsForNodes = new HashMap();
			_templateCountsForNodes.put(node_, templateCountForNode(node_)); //map the node to its count

			visitChildren(node_);
		}
	}

	/**
	 * Visit a node of type {@link FormulaSeq}. Just visit the contained
	 * formulas.
	 */
	public void visit(final FormulaSeq formulaSeq_) {
		final Formulas formulas = formulaSeq_.getFormulas();
		formulas.applyVisitor(this);
	}

	/**
	 * Visits the children of this node.
	 * 
	 * @param node_
	 *            Node to visit
	 */
	private void visitChildren(final Node node_) {
		for (Iterator i = node_.getChildren(); i.hasNext();)
			((Node) i.next()).applyVisitor(this);
	}

	/**
	 * Returns the name of the file containing node_ if file exists. Otherwise,
	 * returns null.
	 */
	private String getLocationFileName(final Node node_) {
		final Location location = node_.getLocation();
		if (location == null || location.getFileName() == null || location.getFileName() == Location.NOT_FROM_FILE)
			return null;
		return location.getFileName();
	}
}
